from flask import Flask, request, jsonify, render_template, session
from flask_cors import CORS
import scipy.io as scio
import numpy as np
import matplotlib.pyplot as plt
import io
import sys
import base64
from scipy.fftpack import fft
import os
import tempfile
import matplotlib
import paddle
from paddle.optimizer import Adam
from paddle import nn

# 获取当前文件的目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 添加上级目录的模块路径到sys.path
sys.path.append(os.path.join(current_dir, "model"))

from main import invoke_model


matplotlib.use("Agg")  # 使用无界面的后端

app = Flask(__name__)
app.secret_key = "123456"  # 设置一个密钥用于会话加密
CORS(app)  # 添加 CORS 支持

# 确保图像保存目录存在
IMAGE_DIR = "images"
if not os.path.exists(IMAGE_DIR):
    os.makedirs(IMAGE_DIR)


@app.route("/")
def index():
    return render_template("index.html")

# 数据处理
@app.route("/upload", methods=["POST"])
def upload_file():
    file = request.files["file"]
    if file and file.filename.endswith(".mat"):
        # 创建临时文件
        temp_file = tempfile.NamedTemporaryFile(delete=False)
        file.save(temp_file.name)

        try:
            # 读取文件内容后立即关闭文件
            data = scio.loadmat(temp_file.name)
            temp_file.close()

            # 检查读取的数据
            if "label" not in data or "data1" not in data or "phi1" not in data:
                raise ValueError(
                    "MAT file does not contain required variables: 'label', 'data1', 'phi1'"
                )

            label = data["label"]
            y = data["data1"]
            phi = data["phi1"]

            label_array = np.array(label)
            y_array = np.array(y)
            phi_array = np.array(phi)

            x_data = []
            for i in range(len(label_array)):
                x_data.append(np.dot(np.linalg.pinv(phi_array[i]), y_array[i]))

            x_data_array = np.array(x_data)
            x_data_array_fft = fft(x_data_array)

            x_data_real = np.real(x_data_array_fft)
            x_data_img = np.imag(x_data_array_fft)

            x_final = np.stack((x_data_real, x_data_img), axis=3)

            images = []

            # 选择一个切片进行可视化（假设是第一个样本的频谱数据）
            x_real = x_final[0, :, :, 0]  # 实部
            x_imag = x_final[0, :, :, 1]  # 虚部

            # 定义时间轴（根据数据情况）
            L = 195
            R = 1
            K = 91
            K0 = 10
            fnyq = 10e10
            TimeResolution = 1 / fnyq
            TimeWin = [0, L * R * K - 1, L * R * (K + K0) - 1]
            t_axis = np.arange(TimeWin[0], TimeWin[-1] + 1) * TimeResolution

            # 定义数字时间轴
            Digital_time_axis = np.linspace(t_axis[0], t_axis[-1], x_real.shape[0])

            # 要可视化的样本数据（选择一个特定的频率分量，例如第一个）
            DigitalSamples1 = x_real[:, 0]
            DigitalSignalSamples = x_imag[:, 0]

            # 绘制 DigitalSamples1
            time_axis_min = 0
            time_axis_max = 2e-7

            # 振幅范围可以手动设置为较小的范围
            amplitude_min = -30000
            amplitude_max = 30000

            # 绘图
            plt.figure(figsize=(10, 6))  # 调整图像大小
            plt.plot(Digital_time_axis, DigitalSamples1, "r")
            plt.title("DigitalSamples1")
            plt.xlabel("times (s)")
            plt.ylabel("magnitude")
            plt.xlim(time_axis_min, time_axis_max)
            plt.ylim(amplitude_min, amplitude_max)
            plt.grid(True)
            plt.tight_layout(pad=3.0)  # 确保坐标轴标题不被截断
            plt.subplots_adjust(right=0.95)  # 调整右边距

            # 将图像保存为 base64 编码的字符串
            img = io.BytesIO()
            plt.savefig(
                img, format="png", bbox_inches="tight"
            )  # 使用 bbox_inches='tight' 确保图像不被截断
            img.seek(0)
            img_base64 = base64.b64encode(img.read()).decode("utf-8")
            images.append(img_base64)
            plt.close()  # 关闭当前图像，防止内存泄漏

            return jsonify({"label": label.tolist(), "images": images})
        except Exception as e:
            return jsonify({"error": str(e)}), 500
        finally:
            # 保存处理后的数据
            data_out = {"x": x_final, "label": label}
            # print("处理后的数据:", data_out)
            temp_mat_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mat")
            scio.savemat(temp_mat_file.name, data_out)
            temp_mat_file.close()
            session["temp_mat_file"] = temp_mat_file.name  # 将临时文件路径存储在会话中
            print("模型测试文件路径:", temp_mat_file.name)
            if not temp_file.closed:
                temp_file.close()
            os.remove(temp_file.name)
    else:
        return jsonify({"error": "Invalid file format"}), 400

# 传入数据进模型进行预测
@app.route("/predict", methods=["POST"])
def predict():
    if "temp_mat_file" in session:
        temp_mat_file_path = session["temp_mat_file"]  # 从会话中获取临时文件路径
        # print("模型测试文件路径:", temp_mat_file_path)
        try:
            # 测试模型并获取准确率
            accuracy = invoke_model(temp_mat_file_path)
            print("模型测试完成，准确率:", accuracy)

            # 返回准确率
            return jsonify({"accuracy": accuracy})

        except Exception as e:
            print("处理.mat文件时出错:", str(e))
            return jsonify({"error": str(e)}), 500
        finally:
            if os.path.exists(temp_mat_file_path):
                os.remove(temp_mat_file_path)  # 删除临时文件
    else:
        return jsonify({"error": "No data available"}), 400


if __name__ == "__main__":
    app.run(debug=True)
